//
//  WaterfallLayout.swift
//  RKTemplate
//
//  Created by zhaoke on 2023/9/25.
//

import UIKit
import Foundation

protocol WaterfallLayoutDataSource: AnyObject {
    func numberOfCols(_ waterfall: WaterfallLayout) -> Int
    func waterfall(_ waterfall: WaterfallLayout, item: Int) -> CGFloat
}

class WaterfallLayout: UICollectionViewFlowLayout {
    weak var dataSource: WaterfallLayoutDataSource?
    fileprivate lazy var cellAttrs: [UICollectionViewLayoutAttributes] = [UICollectionViewLayoutAttributes]()
    fileprivate lazy var cols: Int = {
        return self.dataSource?.numberOfCols(self) ?? 2
    }()
    fileprivate lazy var totalHeights: [CGFloat] = Array.init(repeating: self.sectionInset.top, count: self.cols)
}
extension WaterfallLayout {
    override func prepare() {
        super.prepare()
        let itemCount = collectionView!.numberOfItems(inSection: 0)
        let cellWidth: CGFloat = (collectionView!.bounds.width - sectionInset.left - sectionInset.right - CGFloat(cols - 1) * minimumInteritemSpacing) / CGFloat(cols)
        for i in cellAttrs.count ..< itemCount {
            let indexPath = IndexPath(item: i, section: 0)
            let attr = UICollectionViewLayoutAttributes(forCellWith: indexPath)
            guard let cellHeight: CGFloat = dataSource?.waterfall(self, item: i) else {
                fatalError("请实现对应的数据源方法,并返回Cell的高度")
            }
            let minHeight = totalHeights.min()!
            let minIndex = totalHeights.firstIndex(of: minHeight)!
            let cellX: CGFloat = sectionInset.left + (minimumInteritemSpacing + cellWidth) * CGFloat(minIndex)
            let cellY: CGFloat = minHeight
            attr.frame = CGRect(x: cellX, y: cellY, width: cellWidth, height: cellHeight)
            cellAttrs.append(attr)
            totalHeights[minIndex] = minHeight + minimumLineSpacing + cellHeight
        }
    }
}
extension WaterfallLayout {
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return cellAttrs
    }
    override var collectionViewContentSize: CGSize {
        return CGSize.init(width: 0, height: totalHeights.max()! + sectionInset.bottom - minimumLineSpacing)
    }
}
